import numpy as np
import platgo as pg
import scipy.io as sio

def schaffer(X: np.ndarray) -> np.ndarray:
    X = X ** 2
    return np.sum(0.5 + (np.sin(np.sqrt(X + np.hstack((X[:, 1:], X[:, 0:1])))) ** 2 - 0.5) / (1 + 0.001 * (X +
            np.hstack((X[:, 1:], X[:, 0:1])))) ** 2, axis=1)

def HGBat(X: np.ndarray) -> np.ndarray:
    X = 0.05 * X - 1
    return np.sqrt(np.abs(np.sum(X ** 2, axis=1) ** 2 - np.sum(X, axis=1) ** 2)) + (0.5 + np.sum(X ** 2, axis=1) +
            np.sum(X, axis=1)) / X.shape[1] + 0.5

def Rosenbrock(X: np.ndarray) -> np.ndarray:
    X = 0.02048 * X + 1
    return np.sum(100 * (X[:, 0:X.shape[1] - 1] ** 2 - X[:, 1:]) ** 2 + (X[:, 0:X.shape[1] - 1] - 1) ** 2, axis=1)

def Schwefel(X: np.ndarray) -> np.ndarray:
    X = 10 * X + (4.2097e+2)
    g = X * np.sin(np.sqrt(np.abs(X)))
    temp = 500 - np.mod(X[X > 500], 500)
    g[X > 500] = temp * np.sin(np.sqrt(np.abs(temp))) - (X[X > 500] - 500) ** 2 / 10000 / X.shape[1]
    temp = np.mod(np.abs(X[X < -500]), 500) - 500
    g[X < -500] = temp * np.sin(np.sqrt(np.abs(temp))) - (X[X < -500] - 500) / 10000 / X.shape[1]
    return 418.9829 * X.shape[1] - np.sum(g, axis=1)

def eliptic(X: np.ndarray) -> np.ndarray:
    return np.sum((1e+6) ** (np.tile(np.arange(X.shape[1]), (X.shape[0], 1)) / (X.shape[1] - 1 + (1e-6))) * X ** 2, axis=1)


class CEC_2020_F7(pg.Problem):

    def __init__(self, D=None) -> None:
        self.name = 'CEC_2020_F7'
        self.type['single'], self.type['real'] = [True] * 2
        self.M = 1
        load_path = 'CEC2020.mat'
        load_data = sio.loadmat(load_path)
        mat = []
        for k in load_data.items():
            mat.append(k)
        self.D = D
        self.O = mat[3][1][0][6][0][0][0]
        if self.D is None or self.D < 10:
            self.D = 5
            self.Mat = mat[3][1][0][6][0][0][1]
            self.S = mat[3][1][0][6][0][0][5]
        elif self.D < 15:
            self.D = 10
            self.Mat = mat[3][1][0][6][0][0][2]
            self.S = mat[3][1][0][6][0][0][6]
        elif self.D < 20:
            self.D = 15
            self.Mat = mat[3][1][0][6][0][0][3]
            self.S = mat[3][1][0][6][0][0][7]
        else:
            self.D = 20
            self.Mat = mat[3][1][0][6][0][0][4]
            self.S = mat[3][1][0][6][0][0][8]
        if self.D == 5:
            p = np.array([1, 1, 1, 1, 1])
        else:
            p = np.ceil(np.array([0.1, 0.2, 0.2, 0.2, 0.3]) * self.D)
            p[0] = self.D - np.sum(p[1:])
        p = np.hstack(([0], np.cumsum(p))).astype(int)
        self.S = [self.S[0][p[i]: p[i + 1]] for i in range(len(p) - 1)]
        lb = [-100] * self.D
        ub = [100] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        Y = pop.decs - np.tile(self.O[0][0: pop.decs.shape[1]], (pop.decs.shape[0], 1))
        Z = np.dot(Y, self.Mat.T)
        pop.objv = 2100 + Schwefel(Z[:, self.S[0] - 1]) + HGBat(Z[:, self.S[1] - 1]) + Rosenbrock(Z[:, self.S[2] - 1]) + \
                   Schwefel(Z[:, self.S[3] - 1]) + eliptic(Z[:, self.S[4] - 1])
        pop.objv = pop.objv.reshape(pop.objv.shape[0], 1)

    def get_optimal(self) -> np.ndarray:
        pass


if __name__ == '__main__':
    problem = CEC_2020_F7()
    alg = pg.algorithms.GA(problem=problem, maxgen=100)
    pop = alg.go(100)
    print(pop)
